import pandas as pd
import tushare as ts
from DataControl.orm.pgsql import pgsql
from datetime import date, timedelta, datetime
import time
"""
该类包括对股票数据的操作，包括从数据库拉去股票数据，将数据从tushare拉去到数据库。
由于访问tushare一分钟只能访问200次，需要访问后停止一分钟
"""
class pullDataFromTushare():

    def __init__(self):
        self.pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
        self.database = pgsql()  # 数据库
        self.today = date.today()  # 今日日期
        self.today_str = self.today.strftime('%Y%m%d')  # 今日日期字符串形式
        self.count = 0 #每分钟最多只能访问300次
        self.allts_code = self.getAllTs_code()

    # 每分钟最多只能访问200次，访问200次后暂停1分钟
    def checkOneMunite(self):
        self.count = self.count + 1
        if self.count == 200 :
            print("wait 1min")
            time.sleep(62)
            self.count = 0

    # 获取所有ts_code
    def getAllTs_code(self):
        data = self.pro.stock_basic(exchange='', list_status='L', fields='ts_code,symbol,name,area,industry,list_date')
        self.checkOneMunite()
        res = data.iloc[:, 0]
       # dapan = self.get_dapanfromtushare()
        #for i in range(dapan.shape[0]):
            #print(dapan.iloc[i, 0])
         #   res.loc[len(res)] = dapan.iloc[i, 0]
        #print(res)
        return res

    # 将turnover_rate从tushare拉去到数据库中。对于所有的股票和对应日期，都会加上turnover_rate一列
    # 由于turnover_rate的接口和获取日线数据的接口不一样，插入turnover_rate只能分开来。
    def getTurnover_ratefromtushre(self):
        flag = True
        cnt = 0
        for ts_code in self.allts_code:
            cnt += 1
          #  print(cnt," of ",self.allts_code.shape[0])
            end_date = self.get_newestday_daily(ts_code=ts_code)
            if end_date == None: continue
            #print(ts_code)
            end_date = end_date.strftime("%Y%m%D")
            data = self.pro.daily_basic(ts_code=ts_code, start_date='20180101', end_date=end_date, fields='turnover_rate,ts_code,trade_date')
            data = data[['turnover_rate','ts_code','trade_date']]
            self.checkOneMunite()
            #assert(data.shape[0] == self.get_count_daily(ts_code=ts_code).iloc[0, 0])
            if data.shape[0] != self.get_count_daily(ts_code=ts_code).iloc[0, 0]:continue
            self.database.insertoverturn_rate(data)

    #插入ts_code股票的指数日线行情，从tushare到数据库
    def getindex_daily_from_tushare(self, ts_code, trade_date = None, start_date = None, end_date = None):
        #ts_code：指数代码， trade_date：交易日期 start_date：结束日期 end_date：结束日期
        # 日期格式：YYYYMMDD

        # 有两个接口daily和index_daily，如果一个接口中没有找到对应的股票，去另外一个找
        data = self.pro.daily(ts_code = ts_code, trade_date = trade_date, start_date = start_date, end_date = end_date)
        if data.shape[0] == 0:
            data = self.pro.index_daily(ts_code = ts_code, trade_date = trade_date, start_date = start_date, end_date = end_date)

        #从另外一个接口获取turnover_rate，这项指标不在前面两个接口里面
        turnover_rate = self.pro.daily_basic(ts_code=ts_code, start_date=start_date, end_date=end_date, fields='turnover_rate,ts_code,trade_date')
        if turnover_rate.shape[0] == 0:
            turnover_rate = self.pro.index_dailybasic(ts_code=ts_code, start_date = start_date, end_date=end_date, fields='turnover_rate, ts_code, trade_date')

        flag = False
        # 在获取了某只股票的turnover_rate值后需要和日线接口获取的每一天都对应上
        if data.shape[0] != turnover_rate.shape[0]: flag = True
        if flag == True: return
        for i in range(data.shape[0]):
            if data['trade_date'].iloc[i] != turnover_rate['trade_date'].iloc[i]:
                flag = True
                break
        if flag:
            print(ts_code, "     insert failed")
            return

        print(ts_code, "     inserted")
        data['turnover_rate'] = turnover_rate['turnover_rate']
       # print(data['turnover_rate'])
        #print(type(start_date), type(end_date), data.shape[0])
        self.checkOneMunite()
        self.database.insert_index_daily(data)
        return data.shape[0]

    #插入港股的股票指数日线行情
    """
        目前这个接口没有访问权限(hk_daily)
    """
    def getindex_hkdaily_from_tushare(self, ts_code, trade_date = None, start_date=None, end_date = None):
        data = self.pro.hk_daily(ts_code=ts_code, trade_date = trade_date, start_date = start_date, end_date = end_date)
        self.checkOneMunite()
        self.database.insert_index_daily(data)
        return data.shape[0]

    #插入利润表
    def getincome_from_tushare(self, ts_code, start_date = None, end_date = None):
        data = self.pro.income(ts_code = ts_code, start_date = start_date, end_date = end_date)
        self.checkOneMunite()
        self.database.insert_data_bydf(data, 'income')
        return data.shape[0]

    # 插入资产负债表
    def getbalancesheet_from_tushare(self, ts_code, start_date = None, end_date = None):
        data = self.pro.balancesheet(ts_code = ts_code, start_date = start_date, end_date = end_date)
        self.checkOneMunite()
        self.database.insert_data_bydf(data, 'balancesheet')
        return data.shape[0]

    # 插入流水表
    def getcashflow_from_tushare(self, ts_code, start_date = None, end_date = None):
        data = self.pro.cashflow(ts_code = ts_code, start_date = start_date, end_date = end_date)
        self.checkOneMunite()
        self.database.insert_data_bydf(data, 'cashflow')
        return data.shape[0]

    #从数据库获取日线行情
    def getindex_daily_from_database(self, ts_code, start, end):
        #data = self.database.get_index_daily(ts_code = ts_code, start = start, end = end)
        return self.database.get_index_daily(ts_code = ts_code, start = start, end = end)
        #print(data)
        #return data

    # 获取某一只股票所有的日线行情
    def getallindex_daily_from_database(self, ts_code):
        return self.database.get_all_index_daily(ts_code)

    # 获取存在在数据库中的最新的日期
    def get_newestday_daily(self, ts_code):
        return self.database.get_maxdate_daily(ts_code)['latest_date'][0]

    # 获取存在在数据库中的最新的日期
    def get_newestday_cashflow(self, ts_code):
        return self.database.get_maxdate_cashflow(ts_code)['latest_date'][0]

    # 获取存在在数据库中的最新的日期
    def get_newestday_balancesheet(self, ts_code):
        return self.database.get_maxdate_balancesheet(ts_code)['latest_date'][0]

    # 获取存在在数据库中的最新的日期
    def get_newestday_income(self, ts_code):
        return self.database.get_maxdate_income(ts_code)['latest_date'][0]

    # 获取ts_code在数据库中的日线数据条数
    def get_count_daily(self, ts_code):
        return self.database.getCountOfdaily(ts_code).iloc[0, 0]

    # 获取ts_code最后k条日线数据
    def get_lastk_daily(self, ts_code, k):
        return self.database.getLastkOfdaily(ts_code=ts_code, k=k)

    # 获取大盘数据
    def get_dapanfromtushare(self):
        df = pd.DataFrame()
        trade_date = self.today_str
        while df.empty == True:
            df = self.pro.index_dailybasic(trade_date=trade_date, fields='ts_code')
            trade_date = datetime.strptime(trade_date, "%Y%m%d")
            previousday = trade_date - timedelta(days = 1)
            trade_date = previousday.strftime("%Y%m%d")
        return df


    # 更新表
    def update(self, table): # table是一个字符串，表示更新那一张表包括：（index_daily, cashflow, income, balancesheet）
        tot_ts_code = self.getAllTs_code()
        update_fun = None
        if table == 'index_daily' : update_fun = self.getindex_daily_from_tushare
        elif table == 'cashflow' : update_fun = self.getcashflow_from_tushare
        elif table == 'income' : update_fun = self.getincome_from_tushare
        elif table == 'balancesheet' : update_fun = self.getbalancesheet_from_tushare
        newday_fun = None
        if table == 'index_daily' : newday_fun = self.get_newestday_daily
        elif table == 'cashflow' : newday_fun = self.get_newestday_cashflow
        elif table == 'income' : newday_fun = self.get_newestday_income
        elif table == 'balancesheet' : newday_fun = self.get_newestday_balancesheet
        flag = True
        for ts_code in tot_ts_code:
            #print(ts_code)
            #if ts_code != '000001.SH':
                #flag = False
             #   continue
            #if flag: continue
            latest_day = newday_fun(ts_code)

            #如果有新的上市股票但是数据库中没有记录，那么找不到最近的日期，直接使用2018-1-1作为开始日期
            if latest_day == None : latest_day = date(2018, 1, 1)
            latest_dayplus1 = latest_day + timedelta(days=1)
            latest_day_str = latest_dayplus1.strftime('%Y%m%d')
            update_fun(ts_code=ts_code, start_date=latest_day_str, end_date=self.today_str)

#pull = pullDataFromTushare()
#pull.getindex_hkdaily_from_tushare(ts_code = '00700.HK', start_date = '20180101', end_date = '20240628')
if __name__ == '__main__':
    # pull = pullDataFromTushare()
    # ts = pull.getAllTs_code()
    # df = pull.pro.index_dailybasic(**{
    #     "trade_date": "20240614",
    #     "ts_code": "",
    #     "start_date": "",
    #     "end_date": "",
    #     "limit": "",
    #     "offset": ""
    # }, fields=[
    #     "ts_code",
    #     "turnover_rate",
    #     "turnover_rate_f"
    # ])
    # print(df)
    pull = pullDataFromTushare()
    #pull.getTurnover_ratefromtushre()
    pull.update('index_daily')
    #end_date = pull.get_newestday_daily(ts_code='000042.SZ')
    #end_date = end_date.strftime('%Y%m%D')
    #print(pull.pro.index_dailybasic(ts_code = '000042.SZ', start_date = '20180101', end_date = end_date, fields='ts_code,trade_date,turnover_rate,pe').shape[0])